{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1c03ef1e",
   "metadata": {
    "colab_type": "text",
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/pick.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8547169e",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eeMrMI0-1Dhu",
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    DiagramBuilder,\n",
    "    Integrator,\n",
    "    JacobianWrtVariable,\n",
    "    LeafSystem,\n",
    "    Simulator,\n",
    "    StartMeshcat,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook\n",
    "from manipulation.station import MakeHardwareStation, load_scenario"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87a09cfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f741d649",
   "metadata": {
    "colab_type": "text",
    "id": "UZym7Xo7slPU"
   },
   "source": [
    "# Our first end-effector \"controller\"\n",
    "\n",
    "Let's use the pseudo-inverse of the Jacobian to drive the robot around.  To do that, we'll write a very simple system that looks at the current value of $q$, computes $[J^G]^+$, and uses it to command a constant spatial velocity, $V^G$.\n",
    "\n",
    "We'll only run this controller for a short duration.  Constant spatial velocities aren't something that makes sense for a longer simulation!\n",
    "\n",
    "Make sure you try changing $V^G$, and understand how the command relates to the motion of the robot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87f4783b",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WBTFzB5ZuXvP"
   },
   "outputs": [],
   "source": [
    "# We can write a new System by deriving from the LeafSystem class.\n",
    "# There is a little bit of boiler plate, but hopefully this example makes sense.\n",
    "class PseudoInverseController(LeafSystem):\n",
    "    def __init__(self, plant):\n",
    "        LeafSystem.__init__(self)\n",
    "        self._plant = plant\n",
    "        self._plant_context = plant.CreateDefaultContext()\n",
    "        self._iiwa = plant.GetModelInstanceByName(\"iiwa\")\n",
    "        self._G = plant.GetBodyByName(\"body\").body_frame()\n",
    "        self._W = plant.world_frame()\n",
    "\n",
    "        self.DeclareVectorInputPort(\"iiwa.position\", 7)\n",
    "        self.DeclareVectorOutputPort(\"iiwa.velocity\", 7, self.CalcOutput)\n",
    "\n",
    "    def CalcOutput(self, context, output):\n",
    "        q = self.get_input_port().Eval(context)\n",
    "        self._plant.SetPositions(self._plant_context, self._iiwa, q)\n",
    "        J_G = self._plant.CalcJacobianSpatialVelocity(\n",
    "            self._plant_context,\n",
    "            JacobianWrtVariable.kQDot,\n",
    "            self._G,\n",
    "            [0, 0, 0],\n",
    "            self._W,\n",
    "            self._W,\n",
    "        )\n",
    "        J_G = J_G[:, 0:7]  # Ignore gripper terms\n",
    "\n",
    "        V_G_desired = np.array(\n",
    "            [\n",
    "                0,  # rotation about x\n",
    "                -0.1,  # rotation about y\n",
    "                0,  # rotation about z\n",
    "                0,  # x\n",
    "                -0.05,  # y\n",
    "                -0.1,\n",
    "            ]\n",
    "        )  # z\n",
    "        v = np.linalg.pinv(J_G).dot(V_G_desired)\n",
    "        output.SetFromVector(v)\n",
    "\n",
    "\n",
    "scenario_data = \"\"\"\n",
    "directives:\n",
    "- add_directives:\n",
    "    file: package://manipulation/iiwa_and_wsg.dmd.yaml\n",
    "model_drivers:\n",
    "    iiwa: !IiwaDriver\n",
    "      hand_model_name: wsg\n",
    "    wsg: !SchunkWsgDriver {}\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "def jacobian_controller_example():\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    scenario = load_scenario(data=scenario_data)\n",
    "    station = builder.AddSystem(MakeHardwareStation(scenario, meshcat=meshcat))\n",
    "    plant = station.GetSubsystemByName(\"plant\")\n",
    "\n",
    "    controller = builder.AddSystem(PseudoInverseController(plant))\n",
    "    integrator = builder.AddSystem(Integrator(7))\n",
    "\n",
    "    builder.Connect(controller.get_output_port(), integrator.get_input_port())\n",
    "    builder.Connect(\n",
    "        integrator.get_output_port(), station.GetInputPort(\"iiwa.position\")\n",
    "    )\n",
    "    builder.Connect(\n",
    "        station.GetOutputPort(\"iiwa.position_measured\"),\n",
    "        controller.get_input_port(),\n",
    "    )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    simulator = Simulator(diagram)\n",
    "    context = simulator.get_mutable_context()\n",
    "    station_context = station.GetMyContextFromRoot(context)\n",
    "    station.GetInputPort(\"iiwa.feedforward_torque\").FixValue(\n",
    "        station_context, np.zeros((7, 1))\n",
    "    )\n",
    "    station.GetInputPort(\"wsg.position\").FixValue(station_context, [0.1])\n",
    "    integrator.set_integral_value(\n",
    "        integrator.GetMyContextFromRoot(context),\n",
    "        plant.GetPositions(\n",
    "            plant.GetMyContextFromRoot(context),\n",
    "            plant.GetModelInstanceByName(\"iiwa\"),\n",
    "        ),\n",
    "    )\n",
    "\n",
    "    meshcat.StartRecording(set_visualizations_while_recording=False)\n",
    "    simulator.AdvanceTo(5.0 if running_as_notebook else 0.1)\n",
    "    meshcat.PublishRecording()\n",
    "\n",
    "\n",
    "jacobian_controller_example()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
